import argparse
import os
import torch
import numpy as np
import logging


from utils.config import get_train_config
from engine import TrainerFactory


torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True

SEED = 0
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
np.random.seed(SEED)


def setup_logging(logs_dir):
    BASIC_FORMAT = "[%(asctime)s][%(levelname)s] %(message)s"
    DATE_FORMAT = "%Y-%m-%d %H:%M:%S"
    formatter = logging.Formatter(BASIC_FORMAT, DATE_FORMAT)

    console_handler = logging.StreamHandler()
    console_handler.setFormatter(formatter)
    console_handler.setLevel(level=logging.INFO)

    logs_path = os.path.join(logs_dir, 'logs.txt')
    file_handler = logging.FileHandler(logs_path, mode='a', encoding='utf-8')
    file_handler.setFormatter(formatter)
    file_handler.setLevel(level=logging.INFO)

    logging.basicConfig(level=logging.INFO, handlers=[console_handler, file_handler])


def parse_args():
    parser = argparse.ArgumentParser(description="train net")
    parser.add_argument("--config", type=str, help='model name', required=True)
    parser.add_argument("--dataset", type=str, help='dataset name', required=True)
    parser.add_argument("--val-dataset", type=str, help="val dataset name", default=None)
    parser.add_argument("--gpus", type=str, help="use gpus id", default='0')
    parser.add_argument("--resume", type=bool, help="resume from last training", default=True)
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = parse_args()
    config_path = 'configs/config/{}.yaml'.format(args.config)
    config = get_train_config(config_path)

    config['loader']['train']['dataset'] = args.dataset
    config['loader']['test']['dataset'] = args.val_dataset

    config['output_directory'] = os.path.join(config['output_directory'], f'config-{args.config}', f'train-{args.dataset}', f'val-{args.val_dataset}', 'checkpoint')
    config['model']['devices'] = ['cuda:{}'.format(gpu) for gpu in args.gpus.split(',')]

    os.makedirs(config['output_directory'], exist_ok=True)
    setup_logging(config['output_directory'])

    trainer = TrainerFactory.get(config)
    trainer.resume_or_load(resume=args.resume)
    trainer.train()
